import uproot
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


sig_name = {"../../Data/data_4tau/signal/THDM_hA_4tau_mh_mch_ATLAS_pt10_14TeV_01.root": 3.67719}
bkg_name = {
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbar.root": 4452.73,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarww.root": 0.0139,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarz.root": 0.2136,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_ttbarzz.root": 0.0001,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_wwjj.root": 268.291,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_zjj.root": 151060,
    "/home/yancywww/Code/Workplace/Data/data_4tau/bk/Bk_2l2tau_zz.root": 9.296
}


df_sig_list = []
for name, value in sig_name.items():
    df_tmp = uproot.open(name)["datatrain"].arrays(["HT", "MET"], library="pd")
    df_tmp["weight"] = value
    df_sig_list.append(df_tmp)


df_bkg_list = []
for name, value in bkg_name.items():
    df_tmp = uproot.open(name)["datatrain"].arrays(["HT", "MET"], library="pd")
    df_tmp["weight"] = value
    df_bkg_list.append(df_tmp)


df_sig = pd.concat(df_sig_list)
df_bkg = pd.concat(df_bkg_list)


# 计算加权直方图所需的数据
def calculate_weighted_histogram(df, column1, column2):
    data = df[column1] + df[column2]
    weights = df["weight"]
    hist, bin_edges = np.histogram(data, bins=40, range=[0, 600], weights=weights)
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
    bin_width = bin_edges[1] - bin_edges[0]
    hist_normalized = hist / np.sum(hist) / bin_width # 归一化
    return bin_centers, hist_normalized


# 计算背景和信号的加权直方图
bkg_bin_centers, bkg_hist_normalized = calculate_weighted_histogram(df_bkg, "HT", "MET")
sig_bin_centers, sig_hist_normalized = calculate_weighted_histogram(df_sig, "HT", "MET")


# 假设你的画布尺寸为 8 英寸宽，6 英寸高
fig, ax = plt.subplots(figsize=(8, 6))


# 绘制直方图
# 将 ax.bar 替换为 ax.step
ax.step(bkg_bin_centers, bkg_hist_normalized, color='red', label="bkg", linewidth=2)
ax.step(sig_bin_centers, sig_hist_normalized, color='black', label="sig", linewidth=2)

ax.set_xlim([0,600])
ax.set_ylim([0,0.01])


# 设置标题和坐标轴标签
ax.set_xlabel("HT+MET [GeV]", ha='right', x=1.0)
ax.tick_params(axis='x', labelsize=20)
x_label = ax.xaxis.get_label()
x_label.set_fontsize(20)
ax.set_ylabel('Normalized Events/15 GeV')
ax.tick_params(axis='y', labelsize=20)
y_label = ax.yaxis.get_label()
y_label.set_fontsize(20)


# 调整图例样式
leg = plt.legend(fontsize=20, loc=0)
for line in leg.get_lines():
    line.set_linestyle('-')


# 设置坐标轴参数
ax.tick_params(axis='both', which='both', direction='in', top=True, right=True)


left, bottom, width, height = ax.get_position().bounds
# 调整子图顶部位置，使其上移 1cm
bottom = bottom + 0.1
top = bottom + height
left = left + 0.05
fig.subplots_adjust(left=left, bottom=bottom, right=left + width, top=top)


# 显示图形
plt.show()
